import random
import numpy as np
import matplotlib.pyplot as plt
import wandb
import umap
import torch

from sklearn.manifold import TSNE
from io import BytesIO
from PIL import Image
from lightning import Callback


class LogImageECG(Callback):
    def __init__(self, num_samples: int = 5, leads_to_plot=None) -> None:
        super().__init__()
        self.num_samples = num_samples
        self.all_lead_names = ["I", "II", "III", "AVR", "AVL", "AVF", "V1", "V2", "V3", "V4", "V5", "V6"]
        
        # Default to plotting all leads if none specified
        if leads_to_plot is None:
            self.leads_to_plot = list(range(12))
        else:
            self.leads_to_plot = leads_to_plot
        
        self.lead_order = [self.all_lead_names[i] for i in self.leads_to_plot]
        
        # Initialize the table
        self.columns = ["Patient", "Image Orig", "Image Aug"] + [f"ECG Orig {lead}" for lead in self.lead_order] + [f"ECG Aug {lead}" for lead in self.lead_order]
        self.accumulated_data = []
        self.reset_table()

    def reset_table(self):
        self.table = wandb.Table(columns=self.columns)

    def on_train_batch_end(self, trainer, pl_module, outputs, batch, batch_idx, dataloader_idx=0):
        # Only process data from the first batch of each epoch
        if batch_idx == 0:
            batch = {k: v.detach().cpu().squeeze().numpy() for k, v in batch.items()}

            patients = batch["patient_id"]
            images_orig = batch["image_orig"]
            ecgs_orig = batch["ecg_orig"]  # shape [batch_size, 12, time_steps]
            images_aug = batch["image_aug"]
            ecgs_aug = batch["ecg_aug"]
            rpeaks = batch["rpeaks"]

            all_data = []
            for i in range(len(patients)):
                patient_data = [patients[i], rpeaks[i],images_orig[i], images_aug[i]]
                # Append all original ECG data first
                for lead_idx in self.leads_to_plot:
                    patient_data.append(ecgs_orig[i][lead_idx])  # Append original ECG for each lead
                # Then append all augmented ECG data
                for lead_idx in self.leads_to_plot:
                    patient_data.append(ecgs_aug[i][lead_idx])  # Append augmented ECG for each lead
                all_data.append(patient_data)

            # Randomly select num_samples entries to store in accumulated_data
            selected_samples = random.sample(all_data, min(self.num_samples, len(all_data)))
            self.accumulated_data.extend(selected_samples)

    def on_train_epoch_end(self, trainer, pl_module):
        self.reset_table()
        self.log_images_and_ecgs()
        trainer.logger.experiment.log({"data": self.table}, commit=True)

    def log_images_and_ecgs(self):
        for data in self.accumulated_data:
            row_data = [data[0]]  # Patient ID
            row_data.append(wandb.Video(self.create_gif(data[2], duration=100), fps=10, format="gif"))  # Original image
            row_data.append(wandb.Video(self.create_gif(data[3], duration=100), fps=10, format="gif"))  # Augmented image

            # Calculate the starting index of the ECG data in the list
            # 3 initial fields: Patient ID, Image Orig, Image Aug
            start_index_orig = 4
            start_index_aug = start_index_orig + len(self.leads_to_plot)  # Start after all original ECGs

            # Append all original ECG plots
            for j in range(len(self.leads_to_plot)):
                ecg_orig_index = start_index_orig + j
                lead_name = self.lead_order[j]
                ecg_fig_orig = self.plot_ecg_lead(data[ecg_orig_index], lead_name, 'Original')
                row_data.append(wandb.Image(ecg_fig_orig))
                plt.close(ecg_fig_orig)

            # Append all augmented ECG plots
            for j in range(len(self.leads_to_plot)):
                ecg_aug_index = start_index_aug + j
                lead_name = self.lead_order[j]
                ecg_fig_aug = self.plot_ecg_lead(data[ecg_aug_index], lead_name, 'Augmented', data[1])
                row_data.append(wandb.Image(ecg_fig_aug))
                plt.close(ecg_fig_aug)

            # Add the completed row of data to the table
            self.table.add_data(*row_data)

    def create_gif(self, images, duration):
        images = [(np.clip(img, 0, 1) * 255).astype('uint8') for img in images]
        pil_images = [Image.fromarray(img) for img in images]
        gif_buffer = BytesIO()
        pil_images[0].save(gif_buffer, format='GIF', save_all=True, append_images=pil_images[1:], loop=0, duration=duration)
        gif_buffer.seek(0)
        return gif_buffer

    def plot_ecg_lead(self, ecgs, lead_name, title, rpeaks=None):
        """Plotting for ECG leads with ECG-style formatting."""
        fig, ax = plt.subplots(figsize=(16, 6))
        ax.plot(ecgs, color='black', linewidth=1.2)  # Plotting the ECG data
        
        # ECG-style grid and labels
        ax.set_title(f'{title} ECG - Lead {lead_name}', fontsize=18)
        ax.grid(True, which='both', linestyle='-', linewidth=0.5, color='gray', alpha=0.7)
        ax.axhline(0, color='red', linewidth=1.5, linestyle='--')  # Prominent zero line
        
        # Setting y-limits for better view, assuming ECG amplitude in microvolts
        v_min, v_max = ecgs.min(), ecgs.max()
        v_range = max(abs(v_min), abs(v_max)) * 1.1  # Adding 10% padding
        ax.set_ylim([-v_range, v_range])
        
        # Setting ticks
        ax.set_ylabel('Amplitude', fontsize=12)
        ax.set_xlabel('Time', fontsize=12)

        ax.spines['top'].set_visible(False)
        ax.spines['right'].set_visible(False)
        
        # Minor ticks
        ax.minorticks_on()
        ax.tick_params(axis='both', which='both', direction='out', length=6, width=2, colors='black',
                       grid_color='black', grid_alpha=0.5)
        
        # Fine-tuning grid
        ax.xaxis.set_major_locator(plt.MultipleLocator(125))  # Adjust this depending on your time scale
        ax.yaxis.set_major_locator(plt.MultipleLocator(1.0))  # Adjust depending on the amplitude scale
        
        if rpeaks is not None:
            rpeaks = rpeaks[rpeaks != 0]
            peak_values = ecgs[rpeaks]
            ax.scatter(rpeaks, peak_values, color='red', marker='o', s=40, label='R-peaks')
            ax.legend(loc='upper right')

        return fig


class EmbeddingMonitor(Callback):
    def __init__(self, loss_type):
        super().__init__()
        self.emb_storage = None
        self.loss_type = loss_type

    def on_validation_batch_end(self, trainer, pl_module, outputs, batch, batch_idx, dataloader_idx=0):
        if batch_idx == 0:
            if self.loss_type in ['global', 'both']:
                self.emb_storage = {
                    'ecg_global': outputs['learned_ecg_embeddings_global'].detach().cpu().to(torch.float32), # b, d
                    'image_global': outputs['learned_image_embeddings_global'].detach().cpu().to(torch.float32) # b, d
                }
            if self.loss_type in ['local', 'both']:
                if self.emb_storage is None:
                    self.emb_storage = {}
                self.emb_storage.update({
                    'ecg_local': outputs['learned_ecg_embeddings_local'].detach().cpu().to(torch.float32).reshape(-1, outputs['learned_ecg_embeddings_local'].shape[-1]), # b * T // tubelet_size, d
                    'image_local': outputs['learned_image_embeddings_local'].detach().cpu().to(torch.float32).reshape(-1, outputs['learned_image_embeddings_local'].shape[-1]), # b * T // tubelet_size, d
                })

    def on_validation_epoch_end(self, trainer, pl_module):
        embeddings = self.emb_storage
        if self.loss_type == 'global':
            pairs = [('ecg_global', 'image_global')]
        elif self.loss_type == 'local':
            pairs = [('ecg_local', 'image_local')]
        else:
            pairs = [('ecg_local', 'image_local'), ('ecg_global', 'image_global')]

        for ecg_key, image_key in pairs:
            combined_embeddings = torch.cat([embeddings[ecg_key], embeddings[image_key]], dim=0).numpy()

            # t-SNE plot
            tsne_result = TSNE(n_components=2, random_state=42).fit_transform(combined_embeddings)
            fig, ax = plt.subplots(figsize=(10, 5))
            ax.scatter(tsne_result[:embeddings[ecg_key].shape[0], 0], tsne_result[:embeddings[ecg_key].shape[0], 1], c='blue', label='ECG')
            ax.scatter(tsne_result[embeddings[ecg_key].shape[0]:, 0], tsne_result[embeddings[ecg_key].shape[0]:, 1], c='red', label='MRI')
            ax.set_title(f'{ecg_key} vs {image_key} - t-SNE')
            ax.legend()
            plt.close(fig)
            wandb.log({f"{ecg_key}_vs_{image_key}_tSNE": wandb.Image(fig)})

            # UMAP plot
            umap_result = umap.UMAP(n_neighbors=15, n_components=2, metric='euclidean').fit_transform(combined_embeddings)
            fig, ax = plt.subplots(figsize=(10, 5))
            ax.scatter(umap_result[:embeddings[ecg_key].shape[0], 0], umap_result[:embeddings[ecg_key].shape[0], 1], c='blue', label='ECG')
            ax.scatter(umap_result[embeddings[ecg_key].shape[0]:, 0], umap_result[embeddings[ecg_key].shape[0]:, 1], c='red', label='MRI')
            ax.set_title(f'{ecg_key} vs {image_key} - UMAP')
            ax.legend()
            plt.close(fig)
            wandb.log({f"{ecg_key}_vs_{image_key}_UMAP": wandb.Image(fig)})

        self.emb_storage = None
